from langchain_core.documents import Document
from langchain_core.vectorstores import VectorStoreRetriever
# from library.modelio.openai_gpt import get_response
from library.modelio.zhipu_glm import get_response

from library.game_model.zhipu_glm import get_response as game_get_response
from const import *
from session import *
from library.ledvr import get_qa_chain, get_db

qa_chain = get_qa_chain()

app = Flask(__name__)
app.config['SECRET_KEY'] = os.urandom(24)


@app.route('/', methods=['GET', 'POST'])
def index():
    """
    主页
    :return: 主页
    """
    check_session(session)
    return render_template('index.html')


@app.route('/loadHistory', methods=['GET', 'POST'])
def load_messages():
    """
    加载聊天记录
    :return: 聊天记录
    """
    check_session(session)
    success, message = auth(request.headers, session)
    code = 200  # 200表示云端存储了 node.js改写时若云端不存储则返回201
    if not success:
        return {"code": code, "data": [{"role": "web-system", "content": message}]}
    if session.get('user_id') is None:
        messages_history = [
            {"role": "assistant", "content": project_info},
            {"role": "assistant", "content": "#### 当前浏览器会话为首次请求\n"
                                             "#### 请输入已有用户`id`或创建新的用户`id`。\n"
                                             "- 已有用户`id`请在输入框中直接输入\n"
                                             "- 创建新的用户`id`请在输入框中输入`new:xxx`,其中`xxx`为你的自定义id，请牢记\n"
                                             "- 输入`帮助`以获取帮助提示"}
        ]
    else:
        user_info = get_user_info(session.get('user_id'))
        chat_id = user_info['selected_chat_id']
        messages_history = user_info['chats'][chat_id]['messages_history']
        chat_name = user_info['chats'][chat_id]['name']
        logger.warning(f"用户({session.get('user_id')})加载“{chat_name}”对话的聊天记录，共{len(messages_history)}条记录")
    return {"code": code, "data": messages_history, "message": ""}


@app.route('/downloadUserDictFile', methods=['GET', 'POST'])
def download_user_dict_file():
    """
    下载用户字典文件
    :return: 用户字典文件
    """
    check_session(session)
    if request.headers.get("admin-password") is None:
        success, message = auth(request.headers, session)
        if not success:
            return "未授权，无法下载"
        user_id = request.headers.get("user-id")
        if user_id is None:
            return "未绑定用户，无法下载"
        select_user_dict = LRUCache(USER_SAVE_MAX)
        lock.acquire()
        select_user_dict.put(user_id, all_user_dict.get(user_id))
        lock.release()
        # 存储为临时文件再发送出去
        with tempfile.NamedTemporaryFile(suffix='.pkl', delete=False, mode='wb') as temp_file:
            # 将 Python 对象使用 pickle 序列化保存到临时文件中
            pickle.dump(select_user_dict, temp_file)
        response = make_response(send_file(temp_file.name, as_attachment=True))
        response.headers["Content-Disposition"] = f"attachment; filename={user_id}_of_{USER_DICT_FILE}"
        response.call_on_close(lambda: os.remove(temp_file.name))
        return response

    else:
        if request.headers.get("admin-password") != ADMIN_PASSWORD:
            return "管理员密码错误，无法下载"
        response = make_response(send_file(os.path.join(DATA_DIR, USER_DICT_FILE), as_attachment=True))
        response.headers["Content-Disposition"] = f"attachment; filename={USER_DICT_FILE}"
        return response


@app.route('/uploadUserDictFile', methods=['POST'])
def upload_user_dict_file():
    """
    上传用户字典文件 并合并记录
    :return:
    """
    check_session(session)
    file = request.files.get('file')  # 获取上传的文件
    if file:
        if request.headers.get("admin-password") is None:
            success, message = auth(request.headers, session)
            if not success:
                return "未授权，无法合并用户记录"
            user_id = request.headers.get("user-id")
            if user_id is None:
                return "未绑定用户，无法合并用户记录"
            if not file.filename.endswith(".pkl"):
                return "上传文件格式错误，无法合并用户记录"

            # 读取获取的文件
            upload_user_dict = ""
            with tempfile.NamedTemporaryFile(suffix='.pkl', delete=False, mode='wb') as temp_file:
                file.save(temp_file.name)
            # 将 Python 对象使用 pickle 序列化保存到临时文件中
            try:
                with open(temp_file.name, 'rb') as temp_file:
                    upload_user_dict = pickle.load(temp_file)
            except:
                return "上传文件格式错误，无法解析以及合并用户记录"
            finally:
                os.remove(temp_file.name)
            # 判断是否为LRUCache对象
            if not isinstance(upload_user_dict, LRUCache):
                return "上传文件格式错误，无法合并用户记录"
            lock.acquire()
            user_info = all_user_dict.get(user_id)
            lock.release()
            upload_user_info = upload_user_dict.get(user_id)
            if user_info is None or upload_user_info is None:
                return "仅能合并相同用户id的记录，请确保所上传的记录与当前用户id一致"
            backup_user_dict_file()
            for chat_id in upload_user_info['chats'].keys():
                if user_info['chats'].get(chat_id) is None:
                    user_info['chats'][chat_id] = upload_user_info['chats'][chat_id]
                else:
                    new_chat_id = str(uuid.uuid1())
                    user_info['chats'][new_chat_id] = upload_user_info['chats'][chat_id]
            asyncio_run(save_all_user_dict())
            return '个人用户记录合并完成'
        else:
            if request.headers.get("admin-password") != ADMIN_PASSWORD:
                return "管理员密码错误，无法上传用户记录"
            if not file.filename.endswith(".pkl"):
                return "上传文件格式错误，无法上传用户记录"
            # 读取获取的文件
            upload_user_dict = ""
            with tempfile.NamedTemporaryFile(suffix='.pkl', delete=False, mode='wb') as temp_file:
                file.save(temp_file.name)
            # 将 Python 对象使用 pickle 序列化保存到临时文件中
            try:
                with open(temp_file.name, 'rb') as temp_file:
                    upload_user_dict = pickle.load(temp_file)
            except:
                return "上传文件格式错误，无法解析以及合并用户记录"
            finally:
                os.remove(temp_file.name)
            # 判断是否为LRUCache对象
            if not isinstance(upload_user_dict, LRUCache):
                return "上传文件格式错误，无法合并用户记录"
            backup_user_dict_file()
            lock.acquire()
            for user_id in list(upload_user_dict.keys()):
                if all_user_dict.get(user_id) is None:
                    all_user_dict.put(user_id, upload_user_dict.get(user_id))
                else:
                    for chat_id in upload_user_dict.get(user_id)['chats'].keys():
                        if all_user_dict.get(user_id)['chats'].get(chat_id) is None:
                            all_user_dict.get(user_id)['chats'][chat_id] = upload_user_dict.get(user_id)['chats'][
                                chat_id]
                        else:
                            new_chat_id = str(uuid.uuid1())
                            all_user_dict.get(user_id)['chats'][new_chat_id] = upload_user_dict.get(user_id)['chats'][
                                chat_id]
            lock.release()
            asyncio_run(save_all_user_dict())
            return '所有用户记录合并完成'
    else:
        return '文件上传失败'


@app.route('/loadChats', methods=['GET', 'POST'])
def load_chats():
    """
    加载聊天联系人
    :return: 聊天联系人
    """
    check_session(session)
    success, message = auth(request.headers, session)

    if not check_user_bind(session) or not success:
        chats = []
    else:
        user_info = get_user_info(session.get('user_id'))
        chats = []
        chat_ids = []
        # 如果置顶聊天列表不为空
        if "chat_sticky_list" in user_info:
            for chat_id in user_info['chat_sticky_list']:
                if user_info['chats'].get(chat_id) is not None:
                    chat_ids.append(chat_id)

        for chat_id in user_info['chats'].keys():
            if chat_id not in chat_ids:
                chat_ids.append(chat_id)

        user_info['chat_sticky_list'] = chat_ids

        for i, chat_id in enumerate(chat_ids):
            chat_info = user_info['chats'][chat_id]
            if chat_info['chat_with_history']:
                mode = "continuous"
            else:
                mode = "normal"

            if "assistant_prompt" in chat_info:
                assistant_prompt = chat_info['assistant_prompt']
            else:
                assistant_prompt = ""
            if "context_size" not in chat_info:
                chat_info['context_size'] = 5
            if "context_have" not in chat_info:
                chat_info["context_have"] = 1
            chats.append(
                {"id": chat_id, "name": chat_info['name'], "selected": chat_id == user_info['selected_chat_id'],
                 "assistant_prompt": assistant_prompt, "context_size": chat_info['context_size'],
                 "context_have": chat_info["context_have"], "sticky_number": i,
                 "mode": mode, "messages_total": len(user_info['chats'][chat_id]['messages_history'])})
    code = 200  # 200表示云端存储了 node.js改写时若云端不存储则返回201
    return {"code": code, "data": chats, "message": ""}


@app.route('/returnMessage', methods=['GET', 'POST'])
def return_message():
    """
    获取用户发送的消息，调用get_chat_response()获取回复，返回回复，用于更新聊天框
    :return:
    """
    check_session(session)
    request_data = request.get_json()

    success, message = auth(request.headers, session)
    if not success:
        session.clear()

    messages = request_data.get("messages")
    max_tokens = request_data.get("max_tokens")
    model = request_data.get("model")
    temperature = request_data.get("temperature")
    stream = request_data.get("stream")
    continuous_chat = request_data.get("continuous_chat")
    save_message = request_data.get("save_message")

    send_message = messages[-1].get("content")
    send_time = messages[-1].get("send_time")
    display_time = bool(messages[-1].get("display_time"))
    url_redirect = {"url_redirect": "/", "user_id": None}
    if send_message == "帮助":
        return "### 帮助\n" \
               "1. 输入`new:xxx`创建新的用户id\n " \
               "2. 输入`id:your_id`切换到已有用户id，新会话时无需加`id:`进入已有用户\n" \
               "3. 输入`rename_id:xxx`可将当前用户id更改\n" \
               "4. 输入`查余额`可获得余额信息及最近几天使用量\n" \
               "5. 相关设置也可以在设置面板中进行设置\n" \
               "6. 输入`帮助`查看帮助信息"
    if session.get('user_id') is None:  # 如果当前session未绑定用户
        logger.warning("当前会话为首次请求，用户输入:\t" + send_message)
        if send_message.startswith("new:"):
            user_id = send_message.split(":")[1]
            url_redirect["user_id"] = user_id
            if user_id in all_user_dict:
                session['user_id'] = user_id
                return url_redirect
            user_dict = new_user_dict(user_id, send_time)
            lock.acquire()
            all_user_dict.put(user_id, user_dict)  # 默认普通对话
            lock.release()
            logger.warning(f"创建新的用户id:\t{user_id}")
            session['user_id'] = user_id
            url_redirect["user_id"] = user_id
            return url_redirect
        else:
            user_id = send_message
            user_info = get_user_info(user_id)
            if user_info is None:
                logger.warning(f"用户输入的id{user_id}不存在")
                return "用户id不存在，请重新输入或创建新的用户id"
            else:
                session['user_id'] = user_id
                logger.warning(f"切换到已有用户id:\t{user_id}")
                # 重定向到index
                url_redirect["user_id"] = user_id
                return url_redirect
    else:  # 当存在用户id时
        if send_message.startswith("id:"):
            user_id = send_message.split(":")[1].strip()
            user_info = get_user_info(user_id)
            if user_info is None:
                logger.warning(f"用户尝试切换的的id{user_id}不存在")
                return "用户id不存在，请重新输入或创建新的用户id"
            else:
                session['user_id'] = user_id
                url_redirect["user_id"] = user_id
                logger.warning(f"切换到已有用户id:\t{user_id}")
                # 重定向到index
                return url_redirect
        elif send_message.startswith("new:"):
            user_id = send_message.split(":")[1]
            if user_id in all_user_dict:
                return "用户id已存在，请重新输入或切换到已有用户id"
            session['user_id'] = user_id
            url_redirect["user_id"] = user_id
            user_dict = new_user_dict(user_id, send_time)
            lock.acquire()
            all_user_dict.put(user_id, user_dict)
            lock.release()
            logger.warning(f"创建新的用户id:\t{user_id}")
            return url_redirect
        elif send_message.startswith("delete:"):  # 删除用户
            user_id = send_message.split(":")[1]
            if user_id != session.get('user_id'):
                logger.warning(f"用户({session.get('user_id')})尝试删除用户id({user_id})")
                return "只能删除当前会话的用户id"
            else:
                lock.acquire()
                all_user_dict.delete(user_id)
                lock.release()
                session['user_id'] = None
                logger.warning(f"删除用户id:\t{user_id}成功")
                # 异步存储all_user_dict
                asyncio_run(save_all_user_dict())
                return url_redirect
        elif send_message.startswith("set_apikey:"):
            apikey = send_message.split(":")[1]
            user_info = get_user_info(session.get('user_id'))
            user_info['apikey'] = apikey
            # TODO 前端未存储
            logger.info(f"设置用户专属apikey:\t{apikey}")
            return "设置用户专属apikey成功"
        elif send_message.startswith("rename_id:"):
            new_user_id = send_message.split(":")[1]
            user_info = get_user_info(session.get('user_id'))
            if new_user_id in all_user_dict:
                return "用户id已存在，请重新输入"
            else:
                lock.acquire()
                all_user_dict.delete(session['user_id'])
                all_user_dict.put(new_user_id, user_info)
                lock.release()
                session['user_id'] = new_user_id
                asyncio_run(save_all_user_dict())
                logger.warning(f"修改用户id:\t{new_user_id}")
                url_redirect["user_id"] = new_user_id
                return url_redirect
        elif send_message == "查余额":
            user_info = get_user_info(session.get('user_id'))
            apikey = user_info.get('apikey')
            return get_balance(apikey)
        elif '创建赛事' in send_message:
            user_id = session.get('user_id')
            for m in messages:
                keys = list(m.keys())
                for k in keys:
                    if k not in ['role', 'content']:
                        del m[k]
            #根据用户输入提取的字段
            new_filed_dict = game_get_response(False,send_message)
            old_filed_dict = {}
            if user_id not in  all_user_game_dict:
                all_user_game_dict[user_id] = new_filed_dict
            else:
                old_filed_dict = all_user_game_dict[user_id]
            #更新用户赛事创建字段
            if old_filed_dict != {}:
                for key in new_filed_dict:
                    if key == ORGANIZER_List:
                        if new_filed_dict[key][0]["content"] != '':
                            old_filed_dict[key] = new_filed_dict[key]
                    else:
                       if new_filed_dict[key]["content"] != '':
                           old_filed_dict[key] = new_filed_dict[key]
            else:
                old_filed_dict = new_filed_dict
            not_extract_fild_dict = {}
            #封装没有提取到的字段提示用户
            buid_not_extract_fild_dict(old_filed_dict, not_extract_fild_dict)
            if not_extract_fild_dict != {}:
                response = game_get_response(True,json.dumps(not_extract_fild_dict, ensure_ascii=False))
                return response
            else:
                pop = all_user_game_dict.pop(user_id)
                logger.info(f"用户({session.get('user_id')})创建赛事字段:{pop}...")
                # return "{\"success\":\"创建赛事所需信息填充完毕\"}"
                return  json.dumps(pop, ensure_ascii=False)
        else:  # 处理聊天数据
            user_id = session.get('user_id')
            logger.info(f"用户({user_id})发送消息:{send_message}")
            user_info = get_user_info(user_id)
            chat_id = user_info['selected_chat_id']
            messages_history = user_info['chats'][chat_id]['messages_history']
            chat_with_history = user_info['chats'][chat_id]['chat_with_history']
            print(f"messages_history: {json.dumps(messages_history, ensure_ascii=False)}")
            print(f"chat_with_history: {json.dumps(chat_with_history, ensure_ascii=False)}")
            apikey = user_info.get('apikey')
            if chat_with_history:
                user_info['chats'][chat_id]['have_chat_context'] += 1
            if display_time:
                messages_history.append({'role': 'web-system', "content": send_time})
            for m in messages:
                keys = list(m.keys())
                for k in keys:
                    if k not in ['role', 'content']:
                        del m[k]
            if not STREAM_FLAG:
                if save_message:
                    messages_history.append(messages[-1])

                retriever: VectorStoreRetriever = get_db().as_retriever(
                    search_type='similarity',
                    search_kwargs={"k": 5}
                )
                docs: List[Document] = retriever.invoke(send_message)
                response = get_response(messages, docs, send_message)

                # response = qa_chain.invoke({
                #     "query": send_message
                # })['result']

                if save_message:
                    messages_history.append({"role": "assistant", "content": response})
                asyncio_run(save_all_user_dict())

                logger.info(f"用户({session.get('user_id')})得到的回复消息:{response}...")
                # 异步存储all_user_dict
                asyncio_run(save_all_user_dict())
                return response
            else:
                if save_message:
                    messages_history.append(messages[-1])
                asyncio_run(save_all_user_dict())
                if not save_message:
                    messages_history = []
                response = qa_chain.invoke({
                    "query": send_message
                })['result']
                generate = get_response_stream_generate_from_ChatGPT_API(
                    messages,
                    apikey,
                    messages_history,
                    model=model,
                    temperature=temperature,
                    max_tokens=max_tokens
                )
                return app.response_class(generate(),
                                          mimetype='application/json')


def buid_not_extract_fild_dict(filed_dict, not_extract_fild_dict):
    if  filed_dict[ADDRESS]["content"]=='':
        not_extract_fild_dict[ADDRESS] = all_filed_dict[ADDRESS]
    if filed_dict[SITE_ADDRESS]["content"]=='':
        not_extract_fild_dict[SITE_ADDRESS] = all_filed_dict[SITE_ADDRESS]
    if filed_dict[SITES]["content"]=='':
        not_extract_fild_dict[SITES] = all_filed_dict[SITES]
    if filed_dict[START_TIME]["content"]=='':
        not_extract_fild_dict[START_TIME] = all_filed_dict[START_TIME]
    if filed_dict[END_TIME]["content"]=='':
        not_extract_fild_dict[END_TIME] = all_filed_dict[END_TIME]
    if filed_dict[DEADLINE]["content"]=='':
        not_extract_fild_dict[DEADLINE] = all_filed_dict[DEADLINE]
    if filed_dict[ORGANIZER_List][0]["content"]=="":
        not_extract_fild_dict[ORGANIZER_List] = all_filed_dict[ORGANIZER_List]


@app.route('/selectChat', methods=['GET'])
def select_chat():
    """
    选择聊天对象
    :return:
    """
    chat_id = request.args.get("id")
    check_session(session)
    if not check_user_bind(session):
        return {"code": -1, "msg": "请先创建或输入已有用户id"}
    user_id = session.get('user_id')
    user_info = get_user_info(user_id)
    user_info['selected_chat_id'] = chat_id
    return {"code": 200, "msg": "选择聊天对象成功"}


@app.route('/newChat', methods=['GET'])
def new_chat():
    """
    新建聊天对象
    :return:
    """
    name = request.args.get("name")
    time = request.args.get("time")
    new_chat_id = request.args.get("chat_id")
    check_session(session)
    if not check_user_bind(session):
        return {"code": -1, "msg": "请先创建或输入已有用户id"}
    user_id = session.get('user_id')
    user_info = get_user_info(user_id)
    # new_chat_id = str(uuid.uuid1())
    user_info['selected_chat_id'] = new_chat_id
    user_info['chats'][new_chat_id] = new_chat_dict(user_id, name, time)
    user_info["chat_sticky_list"].insert(1, new_chat_id)
    logger.info("新建聊天对象")
    asyncio_run(save_all_user_dict())
    return {
        "code": 200, "data": {
            "name": name,
            "id": new_chat_id,
            "selected": True,

            "messages_total": len(user_info['chats'][new_chat_id]['messages_history'])
        }
    }


@app.route('/deleteHistory', methods=['GET'])
def delete_history():
    """
    清空上下文
    :return:
    """
    check_session(session)
    if not check_user_bind(session):
        logger.info("请先创建或输入已有用户id")
        return {"code": -1, "msg": "请先创建或输入已有用户id"}
    user_info = get_user_info(session.get('user_id'))
    chat_id = user_info['selected_chat_id']
    default_chat_id = user_info['default_chat_id']
    if default_chat_id == chat_id:
        logger.warning("清空历史记录")
        user_info["chats"][chat_id]['messages_history'] = user_info["chats"][chat_id]['messages_history'][:5]
    else:
        logger.warning("删除聊天对话")
        del user_info["chats"][chat_id]
    user_info['selected_chat_id'] = default_chat_id
    return "2"


@app.route('/editChat', methods=['GET', 'POST'])
def edit_chat():
    check_session(session)
    if not check_user_bind(session):
        return {"code": -1, "msg": "请先创建或输入已有用户id"}
    user_id = session.get('user_id')
    user_info = get_user_info(user_id)
    data = request.get_json()
    id = data.get("id")
    if id not in user_info["chats"]:
        return {"code": -1, "msg": "需要编辑的聊天对话不存在"}
    if data.get("name") is not None:
        user_info["chats"][id]["name"] = data.get("name")
    if data.get("context_size") is not None:
        user_info["chats"][id]["context_size"] = data.get("context_size")  # 每次发送请求时，发送的上下文数量
    if data.get("mode") is not None:
        mode = data.get("mode")
        if mode == "normal":
            user_info["chats"][id]['chat_with_history'] = False
        else:
            user_info["chats"][id]['chat_with_history'] = True
    if data.get("assistant_prompt") is not None:
        assistant_prompt = data.get("assistant_prompt")
        user_info["chats"][id]["assistant_prompt"] = assistant_prompt

    if data.get("context_have") is not None:
        context_have = data.get("context_have")
        user_info["chats"][id]["context_have"] = context_have

    # 如果需要置顶该聊天
    if data.get("sticky_number") is not None and id != user_info['default_chat_id']:
        sticky_number = data.get("sticky_number")
        if "chat_sticky_list" not in user_info:
            user_info["chat_sticky_list"] = []
        if sticky_number < 1:
            user_info["chat_sticky_list"].remove(id)
        else:
            if id in user_info["chat_sticky_list"]:
                user_info["chat_sticky_list"].remove(id)
            user_info["chat_sticky_list"].insert(sticky_number, id)

    return {"code": 200, "msg": "修改成功"}


if __name__ == '__main__' or __name__ == 'main':
    logger.warning("持久化存储文件路径为:{}".format(os.path.join(os.getcwd(), os.path.join(DATA_DIR, USER_DICT_FILE))))

    check_load_pickle()

    if len(API_KEY) == 0:
        # 退出程序
        logger.error("请在openai官网注册账号，获取api_key填写至程序内或命令行参数中")
        exit()
    if os.getenv("DEPLOY_ON_ZEABUR") is None:
        app.run(host="0.0.0.0", port=PORT, debug=False)
